import math
import socket
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.table import StreamTableEnvironment, DataTypes
from pyflink.table.descriptors import Schema
import numpy as np
import torch
from transformers import BertTokenizer, BertModel
from Classifier import Classifier
import warnings
warnings.filterwarnings("ignore")

model_name = '../bert-base-chinese'
model_path = '../bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
bert_model = BertModel.from_pretrained(model_path)
model = Classifier(bert_model)
# 加载最佳模型的权重
路径='../文本分类plus/saved_weights_外卖.pt'
model.load_state_dict(torch.load(路径))
good_num = 194
bad_num = 80
class CommentAnalyzer:
    def __init__(self):
        self.env = StreamExecutionEnvironment.get_execution_environment()
        self.t_env = StreamTableEnvironment.create(self.env)

    def process_comments(self):
        global good_num, bad_num
        self.env.set_parallelism(1)

        # 假设使用Socket服务器来接收数据
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.bind(('localhost', 9999))
        server_socket.listen(1)

        while True:
            conn, addr = server_socket.accept()
            with conn:
                data = conn.recv(1024).decode('utf-8').strip()
                # print(f"接收信息: {data}")

                sent_id = tokenizer.encode(str(data),
                                           add_special_tokens=True,
                                           max_length=100,
                                           truncation=True,
                                           pad_to_max_length='right')

                att_mask = [int(tok > 0) for tok in sent_id]
                sent_id = torch.tensor(sent_id)
                att_mask = torch.tensor(att_mask)
                sent_id = sent_id.unsqueeze(0)
                att_mask = att_mask.unsqueeze(0)
                preds = model(sent_id, att_mask)

                total_preds = []
                total_preds.append(np.argmax(preds.detach().cpu().numpy(), axis=1))
                # print(np.concatenate(total_preds, axis=0))
                if 1 in np.concatenate(total_preds, axis=0):
                    print("-------------------------------")
                    print(f"评论: {data}, 评论类型: 好评")
                    good_num+=1
                    print(f"好评: {good_num}, 差评: {bad_num}")
                    result1 = float(good_num) / float(good_num+bad_num)
                    print("好评占比{:.2f}".format(result1))


                else:
                    print("-------------------------------")
                    print(f"评论: {data}, 评论类型: 差评")
                    bad_num+=1
                    print(f"好评: {good_num}, 差评: {bad_num}")
                    result1 = float(good_num) / float(good_num + bad_num)
                    print("好评占比{:.2f}".format(result1))




if __name__ == '__main__':
    analyzer = CommentAnalyzer()
    analyzer.process_comments()
